package spendreport;

import org.apache.flink.api.common.functions.RuntimeContext;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.sink.SinkFunction;
import org.apache.flink.streaming.api.functions.source.SourceFunction;
import org.apache.flink.streaming.connectors.nifi.NiFiDataPacket;
import org.apache.flink.streaming.connectors.nifi.NiFiDataPacketBuilder;
import org.apache.flink.streaming.connectors.nifi.NiFiSink;
import org.apache.flink.streaming.connectors.nifi.NiFiSource;
import org.apache.nifi.remote.client.SiteToSiteClient;
import org.apache.nifi.remote.client.SiteToSiteClientConfig;

import java.nio.charset.StandardCharsets;
import java.util.Map;

public class NifiTest {

    public static void main(String[] args) throws Exception {
        final StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();

        SiteToSiteClientConfig clientConfig = new SiteToSiteClient.Builder()
                .url("http://localhost:9090/nifi/")
                .portName("Flink")
                .requestBatchCount(5)
                .buildConfig();

        SourceFunction<NiFiDataPacket> nifiSource = new NiFiSource(clientConfig);
        DataStream<NiFiDataPacket> stream = env
                .addSource(nifiSource)
                .map(s -> s);
        SinkFunction<NiFiDataPacket> nifiSink = new NiFiSink<>(clientConfig, new NiFiDataPacketBuilder<NiFiDataPacket>() {
            @Override
            public NiFiDataPacket createNiFiDataPacket(NiFiDataPacket niFiDataPacket, RuntimeContext runtimeContext) {
                return new NiFiDataPacket() {
                    @Override
                    public byte[] getContent() {
                        return (new String(niFiDataPacket.getContent()) + " from flink").getBytes(StandardCharsets.UTF_8);
                    }

                    @Override
                    public Map<String, String> getAttributes() {
                        return niFiDataPacket.getAttributes();
                    }
                };
            }
        });
        stream.addSink(nifiSink);
        env.execute();
    }

}